from flask import Flask, jsonify, request, render_template
import cv2
import numpy as np
import base64
import tempfile
import os
import time
import threading
import uuid
import json
from datetime import datetime, timedelta
from standalone import start_stimulus_extraction, start_stimulus_param_user_input

app = Flask(__name__, template_folder='templates')

PORT = 5020
BASE_URL = f"http://127.0.0.1:{PORT}"

# 存储任务状态的全局变量
tasks = {}

class StimulusExtractionTask:
    def __init__(self, video_path, threshold=150):
        self.video_path = video_path
        self.threshold = threshold
        self.progress = 0
        self.status = "Pending"
        self.result = None
        self.point = None  # 用户选择的点
        self.first_frame_base64 = None
        self.thread = None
        self.error = None
    
    def extract_first_frame(self):
        """提取视频第一帧并转换为base64"""
        try:
            # 打开视频
            cap = cv2.VideoCapture(self.video_path)
            if not cap.isOpened():
                raise Exception(f"无法打开视频: {self.video_path}")
            
            # 读取第一帧
            ret, frame = cap.read()
            if not ret:
                raise Exception("无法读取视频第一帧")
            
            # 获取视频尺寸
            self.frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
            self.frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
            
            # 转换为base64
            _, buffer = cv2.imencode('.jpg', frame)
            self.first_frame_base64 = base64.b64encode(buffer).decode('utf-8')
            
            # 关闭视频
            cap.release()
            return self.first_frame_base64
            
        except Exception as e:
            self.error = str(e)
            self.status = "Error"
            print(f"提取第一帧出错: {e}")
            return None
    
    def process_video(self):
        """处理视频并提取刺激帧"""
        try:
            self.status = "Running"
            
            # 打开视频
            cap = cv2.VideoCapture(self.video_path)
            if not cap.isOpened():
                raise Exception(f"无法打开视频: {self.video_path}")
            
            # 获取视频信息
            frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
            
            # 用户选择的点坐标
            x, y = self.point
            
            # 定义判断区域 (7x7)
            domain_top = max(0, y - 3)
            domain_bottom = min(self.frame_height - 1, y + 3)
            domain_left = max(0, x - 3)
            domain_right = min(self.frame_width - 1, x + 3)
            
            # 结果数组
            stimulus_frames = []
            is_stimulus = [False] * frame_count
            
            # 分析每一帧
            for frame_index in range(frame_count):
                # 读取当前帧
                ret, frame = cap.read()
                if not ret:
                    print(f"警告: 无法读取帧 {frame_index}")
                    continue
                
                # 转换为灰度图
                gray_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
                
                # 提取指定区域
                roi = gray_frame[domain_top:domain_bottom+1, domain_left:domain_right+1]
                
                # 手动寻找区域中的最大亮度值 (直接按照light.py的方式实现)
                max_value = 0
                height, width = roi.shape
                for i in range(height):
                    for j in range(width):
                        if roi[i, j] > max_value:
                            max_value = roi[i, j]
                
                # 判断是否为刺激帧
                if max_value > self.threshold:
                    stimulus_frames.append(frame_index)
                    is_stimulus[frame_index] = True
                    print(f"检测到刺激帧: {frame_index + 1}, 最大亮度: {max_value}")
                
                # 更新进度
                self.progress = int((frame_index + 1) / frame_count * 100)
                
                # 定期打印进度信息
                if (frame_index + 1) % 100 == 0 or frame_index == 0 or frame_index == frame_count - 1:
                    print(f"已处理 {frame_index + 1}/{frame_count} 帧 ({self.progress}%)")
            
            # 关闭视频
            cap.release()
            
            # 保存结果
            self.result = stimulus_frames
                # "stimulus_frames": stimulus_frames,
                # "is_stimulus": is_stimulus,
                # "description": "刺激帧索引和每帧状态布尔值"
            
            self.status = "Completed"
            print(f"分析完成，共找到 {len(stimulus_frames)} 个刺激帧")
            
        except Exception as e:
            self.error = str(e)
            self.status = "Error"
            print(f"处理视频出错: {e}")
    
    def run(self):
        """在新线程中运行视频处理"""
        self.thread = threading.Thread(target=self.process_video)
        self.thread.daemon = True
        self.thread.start()

@app.route('/')
def index():
    return 'Stimulus Extraction Service'

@app.route('/meta')
def meta():
    return jsonify({
        'name': 'stimulus-extract',
        'description': '视频刺激提取服务',
        'category': 'video processing',
    })

@app.route('/tasks')
def get_tasks():
    return jsonify([
            'extract_stimulus: (FilePath, ?Int) -> List<Int>',
            'stimulus_param_user_input: Void->(?Bool, Float, ?Float, ?FLoat)'
        ])

@app.route('/workflows')
def get_workflows():
    return jsonify({
        'workflows': []
    })

@app.route('/mount', methods=['GET'])
def mount():
    return jsonify({
        'session_key': str(uuid.uuid4()),
        'healthy_info': 'Healthy'
    })

@app.route('/task/extract_stimulus', methods=['GET', 'POST'])
def start_extraction_task_standalone():
    try:
        # 从请求中获取视频路径和阈值
        data = request.json
        video_path = data[0]
        threshold = int(data[1]) if len(data) > 1 and data[1] else 150
        
        # 创建任务ID
        task_id = str(uuid.uuid4())
        
        # 创建任务实例
        task = StimulusExtractionTask(video_path, threshold)
        
        # 提取第一帧
        first_frame = task.extract_first_frame()
        if not first_frame:
            return jsonify({"error": task.error}), 400
        
        # 存储任务
        tasks[task_id] = task

        task_info = {
            'task_id': task_id,
            'ui_addr': f'/task/extract_stimulus/ui/{task_id}',
            'progress_pipe_addr': f'/task/extract_stimulus/progress/{task_id}',
            'result_addr': f'/task/extract_stimulus/result/{task_id}'
        }

        threading.Thread(target=start_stimulus_extraction, args=(task_info, video_path, threshold)).start()
        
        # task_info['ui_addr'] = BASE_URL + task_info['ui_addr']
        # task_info['progress_pipe_addr'] = BASE_URL + task_info['progress_pipe_addr']
        # task_info['result_addr'] = BASE_URL + task_info['result_addr']

        # 返回任务信息和UI URL
        return jsonify(task_info)
        
    except Exception as e:
        return jsonify({"error": str(e)}), 400
    
@app.route('/task/0/extract_stimulus', methods=['POST'])
def start_extraction_task():
    try:
        # 从请求中获取视频路径和阈值
        data = request.json
        video_path = data[0]
        threshold = int(data[1]) if len(data) > 1 else 150
        
        # 创建任务ID
        task_id = str(uuid.uuid4())
        
        # 创建任务实例
        task = StimulusExtractionTask(video_path, threshold)
        
        # 提取第一帧
        first_frame = task.extract_first_frame()
        if not first_frame:
            return jsonify({"error": task.error}), 400
        
        # 存储任务
        tasks[task_id] = task
        
        # 返回任务信息和UI URL
        return jsonify({
            'task_id': task_id,
            'ui_addr': f'/task/extract_stimulus/ui/{task_id}',
            'progress_pipe_addr': f'/task/extract_stimulus/progress/{task_id}',
            'result_addr': f'/task/extract_stimulus/result/{task_id}'
        })
        
    except Exception as e:
        return jsonify({"error": str(e)}), 400

@app.route('/task/stimulus_param_user_input', methods=['POST'])
def stimulus_param_user_input_standalone():
    task_id = str(uuid.uuid4())
    
    # 创建任务对象并存储在全局任务字典中
    tasks[task_id] = {
        'status': 'Pending',
        'progress': 0,
        'error': None,
        'result': None
    }
    
    task_info = {
        'task_id': task_id,
        'ui_addr': f'/task/stimulus_param_user_input/ui/{task_id}',
        'progress_pipe_addr': f'/task/stimulus_param_user_input/progress/{task_id}',
        'result_addr': f'/task/stimulus_param_user_input/result/{task_id}'
    }

    threading.Thread(target=start_stimulus_param_user_input, args=(task_info,)).start()
    
    print(f"创建刺激参数输入任务: {task_id}")
    return jsonify(task_info)

@app.route('/task/0/stimulus_param_user_input', methods=['GET'])
def stimulus_param_user_input():
    task_id = str(uuid.uuid4())
    
    # 创建任务对象并存储在全局任务字典中
    tasks[task_id] = {
        'status': 'Pending',
        'progress': 0,
        'error': None,
        'result': None
    }
    
    task_info = {
        'task_id': task_id,
        'ui_addr': f'/task/stimulus_param_user_input/ui/{task_id}',
        'progress_pipe_addr': f'/task/stimulus_param_user_input/progress/{task_id}',
        'result_addr': f'/task/stimulus_param_user_input/result/{task_id}'
    }
    
    print(f"创建刺激参数输入任务: {task_id}")
    return jsonify(task_info)

@app.route('/task/stimulus_param_user_input/result_input/<task_id>', methods=['POST'])
def stimulus_param_result_input(task_id):
    try:
        data = request.json
        print("收到刺激参数:", data)
        
        # 验证数据
        if not data.get('position'):
            return jsonify({"error": "缺少刺激位置参数"}), 400
            
        if not data.get('frequency'):
            return jsonify({"error": "缺少刺激频率参数"}), 400
            
        if not data.get('voltage') and not data.get('current'):
            return jsonify({"error": "刺激电压和电流至少提供一个"}), 400
        
        task = tasks[task_id]
        task['status'] = 'Completed'
        task['progress'] = 100
        task['result'] = data
        
        # 保存数据到文件
        # with open('stimulus_params.json', 'w', encoding='utf-8') as f:
            # json.dump(data, f, ensure_ascii=False, indent=2)
            
        return jsonify({"status": "success", "message": "参数已保存"})
    except Exception as e:
        print(f"处理刺激参数出错: {e}")
        task['status'] = 'Failed'
        task['error'] = str(e)
        return jsonify({"error": str(e)}), 400

@app.route('/task/stimulus_param_user_input/ui/<task_id>')
def stimulus_param_form(task_id):
    """显示刺激参数输入表单"""
    if task_id not in tasks:
        return jsonify({"error": "任务不存在"}), 404
    
    task = tasks[task_id]
    return render_template('stimulus_params.html', task_id=task_id)

@app.route('/task/stimulus_param_user_input/result/<task_id>')
def get_stimulus_params(task_id):
    """获取保存的刺激参数"""

    if task_id not in tasks:
        return jsonify({"error": "任务不存在"}), 404
    
    task = tasks[task_id]
    params = task.get('result') if isinstance(task, dict) else task.result
    
    if params is None:
        return jsonify({"error": "尚未提交参数"}), 404
    
    try:
        return jsonify(params)
    except Exception as e:
        return jsonify({"error": str(e)}), 500

@app.route('/task/stimulus_param_user_input/progress/<task_id>')
def stimulus_param_result_progress(task_id):
    if task_id not in tasks:
        return jsonify({"error": "任务不存在"}), 404
    
    task = tasks[task_id]
    
    # 根据task类型选择正确的访问方式
    if isinstance(task, dict):
        return jsonify({
            'status': task.get('status', 'Pending'),
            'percentage': task.get('progress', 0),
            'error': task.get('error')
        })
    else:
        return jsonify({
            'status': task.status,
            'percentage': task.progress,
            'error': task.error
        })

@app.route('/task/extract_stimulus/ui/<task_id>')
def extraction_ui(task_id):
    if task_id not in tasks:
        return "任务不存在", 404
    
    task = tasks[task_id]
    return render_template('extraction.html', 
                          task_id=task_id, 
                          image_data=task.first_frame_base64,
                          video_path=task.video_path)

@app.route('/task/extract_stimulus/select_point/<task_id>', methods=['POST'])
def select_point(task_id):
    if task_id not in tasks:
        return jsonify({"error": "任务不存在"}), 404
    
    task = tasks[task_id]
    
    try:
        data = request.json
        x = int(data['x'])
        y = int(data['y'])
        
        # 获取用户自定义阈值
        threshold = int(data.get('threshold', 150))
        
        print(f"收到选择点坐标: ({x}, {y}), 图像尺寸: {task.frame_width}x{task.frame_height}, 阈值: {threshold}")
        
        # 确保坐标在有效范围内
        x = max(0, min(x, task.frame_width - 1))
        y = max(0, min(y, task.frame_height - 1))
        
        print(f"调整后的坐标: ({x}, {y})")
        
        # 保存选择的点
        task.point = (x, y)
        
        # 更新阈值
        task.threshold = threshold
        
        # 计算判断区域并打印
        domain_top = max(0, y - 3)
        domain_bottom = min(task.frame_height - 1, y + 3)
        domain_left = max(0, x - 3)
        domain_right = min(task.frame_width - 1, x + 3)
        
        print(f"判断区域: 上={domain_top}, 下={domain_bottom}, 左={domain_left}, 右={domain_right}")
        print(f"使用亮度阈值: {threshold}")
        
        # 开始处理
        task.run()
        
        return jsonify({"status": "success"})
        
    except Exception as e:
        print(f"选择点错误: {e}")
        return jsonify({"error": str(e)}), 400

@app.route('/task/extract_stimulus/progress/<task_id>')
def get_progress(task_id):
    if task_id not in tasks:
        return jsonify({"error": "任务不存在"}), 404
    
    task = tasks[task_id]
    
    return jsonify({
        'status': task.status,
        'percentage': task.progress,
        'error': task.error
    })

@app.route('/task/extract_stimulus/result/<task_id>')
def get_result(task_id):
    if task_id not in tasks:
        return jsonify({"error": "任务不存在"}), 404
    
    task = tasks[task_id]
    
    if task.status != "Completed":
        return jsonify({"error": "任务尚未完成"}), 400
    
    return jsonify(task.result)

if __name__ == '__main__':
    app.run(host='127.0.0.1', port=PORT, debug=True) 